# -*- coding: utf-8 -*-


class Event:
    def __init__(self, user, item):
        self.user = user
        self.item = item


class Data_set:
    def __init__(self):
        self.user_events = {}  # 与user交互过的所有item
        self.item_events = {}  # 与item交互过的所有user
        self.events = []  # 所有事件集，测试集才有

    def data_load(self, path: str, phase):
        """
            训练数据读取
        :param phase: 标志训练集和测试集
        :param path: 读取文件路径
        :return:
        """
        with open(path, 'r') as train:
            for line in train.readlines():
                user, item = line.split()[:2]
                if phase == 'train':
                    if user not in self.user_events:
                        self.user_events[user] = set()
                    if item not in self.item_events:
                        self.item_events[item] = set()
                    self.user_events[user].add(item)
                    self.item_events[item].add(user)
                if phase == 'test':
                    self.events.append(Event(user, item))


path_train = r'../data_MF/train.txt'
path_test = r'../data_MF/test.txt'

# 训练集
train_data = Data_set()
train_data.data_load(path_train, 'train')

# 测试集
test_data = Data_set()
test_data.data_load(path_test, 'test')


def train_and_test(train_rate):
    """
        数据集划分为训练集和测试集
    :param train_rate:
    :return:
    """
    path_train = r'../data_MF/train.txt'
    path_test = r'../data_MF/test.txt'
    path_all = r'../data_MF/data_all.txt'

    train = open(path_train, 'a')
    train.truncate(0)
    test = open(path_test, 'a')
    test.truncate(0)
    all_data = open(path_all)
    num_all = len(all_data.readlines())
    all_data.close()
    all_data = open(path_all)
    train_num = int(num_all * train_rate)
    test_num = num_all - train_num
    event_idx = 0
    for line in all_data.readlines():
        if event_idx < train_num:
            train.write(line)
        else:
            test.write(line)
        event_idx += 1

    train.close()
    test.close()
    all_data.close()


if __name__ == '__main__':
    train_and_test(0.9)
